import argparse
import json
import os
import sys
from pathlib import Path

from PIL import Image
from tqdm import tqdm

import utils


class Labelbox(object):
    def __init__(self, images: Path, annotation: Path, classes: Path):
        assert all(isinstance(i, Path) for i in (images, annotation)), f'ERROR - {images} or {annotation} is not Path'
        assert os.path.isdir(images) and len(os.listdir(images)), f'{images} must exist and be not empty'
        assert os.path.isfile(annotation), f'{annotation} must exist as a file'
        self.images = images
        self.classes = classes

        with open(annotation, 'r', encoding='utf8') as f:
            self.annotation = json.load(f)

    def impl(self) -> None:
        self.check()
        utils.create_folder(Path('.') / 'output')
        classes_li = self.read_classes()

        with tqdm(total=len(self.annotation), file=sys.stdout) as pbar:
            for i in self.annotation:
                name = i['External ID']
                objects = i['Label']['objects']
                with Image.open(self.images / name) as img, open(Path('.') / 'output' /
                                                                 (os.path.splitext(name)[0] + '.txt'), 'w',
                                                                 encoding='utf8') as f:
                    width, height = img.size
                    for j in range(len(objects)):
                        # calculate relative coords
                        bbox = objects[j]['bbox']  # top, left, height, width
                        coords = [(bbox['left'] + bbox['width'] / 2) / width, (bbox['top'] + bbox['height'] / 2) / height,
                                  bbox['width'] / width, bbox['height'] / height]
                        coords = list(map(lambda x: str(round(x, 6)), coords))
                        cls = objects[j]['value']
                        idx = classes_li.index(cls)

                        # last line handling
                        if j != len(objects) - 1:
                            line = str(idx) + ' ' + ' '.join(coords) + '\n'
                        else:
                            line = str(idx) + ' ' + ' '.join(coords)
                        f.write(line)

                pbar.update(1)

        print('Done!')

    def read_classes(self) -> list:
        assert os.path.isfile(self.classes) and os.path.splitext(self.classes)[1] == '.txt', \
            f'ERROR - {self.classes} invalid'
        with open(self.classes, 'r', encoding='utf8') as f:
            lines = f.readlines()
            return [line.split()[0] for line in lines]

    def check(self) -> None:
        names = []

        with tqdm(total=len(self.annotation), file=sys.stdout) as pbar:
            for i in self.annotation:
                names.append(i['External ID'])
                pbar.update(1)

        if set(os.listdir(self.images)) - set(names):
            raise FileNotFoundError('image names unmatched with names in annotation file')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--images', type=str, default='dataset/images', help='directory of images')
    parser.add_argument('--ann', type=str, required=True, help='annotation file')
    parser.add_argument('--classes', type=str, required=True, help='classes.txt')

    opt = parser.parse_args()
    print(opt)

    labelbox = Labelbox(Path(opt.images), Path(opt.ann), Path(opt.classes))
    labelbox.impl()
